import cv2
import numpy as np
import matplotlib.pyplot as plt


# DCT嵌入类
class DCT_Embed(object):
    def __init__(self, background, watermark, block_size=8, alpha=30):
        np.random.seed(2024)
        b_h, b_w = background.shape[:2]
        assert watermark.size <= (b_h // block_size) * (b_w // block_size), "水印过大，无法嵌入"

        # 参数保存
        self.block_size = block_size
        self.alpha = alpha
        self.k1 = np.random.randn(block_size)
        self.k2 = np.random.randn(block_size)

    def dct_blkproc(self, background):
        background_dct_blocks_h = background.shape[0] // self.block_size  # 高度
        background_dct_blocks_w = background.shape[1] // self.block_size  # 宽度
        background_dct_blocks = np.zeros(shape=((background_dct_blocks_h, background_dct_blocks_w, self.block_size, self.block_size)))

        h_data = np.vsplit(background, background_dct_blocks_h)
        for h in range(background_dct_blocks_h):
            block_data = np.hsplit(h_data[h], background_dct_blocks_w)
            for w in range(background_dct_blocks_w):
                a_block = block_data[w]
                background_dct_blocks[h, w, ...] = cv2.dct(a_block.astype(np.float64))
        return background_dct_blocks

    def dct_embed(self, dct_data, watermark):
        temp = watermark.flatten()
        result = dct_data.copy()
        for h in range(watermark.shape[0]):
            for w in range(watermark.shape[1]):
                k = self.k1 if watermark[h, w] == 1 else self.k2
                for i in range(self.block_size):
                    result[h, w, i, self.block_size - 1] = dct_data[h, w, i, self.block_size - 1] + self.alpha * k[i]
        return result

    def idct_embed(self, dct_data):
        row = None
        result = None
        h, w = dct_data.shape[0], dct_data.shape[1]
        for i in range(h):
            for j in range(w):
                block = cv2.idct(dct_data[i, j, ...])
                row = block if j == 0 else np.hstack((row, block))
            result = row if i == 0 else np.vstack((result, row))
        return result.astype(np.uint8)

    def dct_extract(self, synthesis, watermark_size):
        self.k1 = np.random.randn(self.block_size)
        self.k2 = np.random.randn(self.block_size)

        w_h, w_w = watermark_size
        recover_watermark = np.zeros(shape=watermark_size)
        synthesis_dct_blocks = self.dct_blkproc(background=synthesis)
        p = np.zeros(8)
        for h in range(w_h):
            for w in range(w_w):
                for k in range(self.block_size):
                    p[k] = synthesis_dct_blocks[h, w, k, self.block_size - 1]
                if corr2(p, self.k1) > corr2(p, self.k2):
                    recover_watermark[h, w] = 1
                else:
                    recover_watermark[h, w] = 0
        return recover_watermark


# 计算均值
def mean2(x):
    return np.sum(x) / np.size(x)


# 计算相关性
# def corr2(a, b):
#     a = a - mean2(a)
#     b = b - mean2(b)
#     r = (a * b).sum() / np.sqrt((a * a).sum() * (b * b).sum())
#     return r
def corr2(a, b):
    a = a - mean2(a)
    b = b - mean2(b)
    denominator = np.sqrt((a * a).sum() * (b * b).sum())
    # 确保分母不为零
    if denominator == 0:
        return 0  # 或者返回一个其他的合适的默认值，例如 0
    r = (a * b).sum() / denominator
    return r


# 将文本转换为二进制字符串
def text_to_binary(text):
    return "".join(format(ord(c), "08b") for c in text)

# rsa修改
# 将rsabytes转为二进制字符串
def bytes_to_binary(rsa_bytes):
    binary_list = []
    for byte in rsa_bytes:
        binary_list.append(bin(byte).replace('0b', '').zfill(8))
    return binary_list

# 将二进制转换为文本
def binary_to_text(binary):
    text = ""
    for i in range(0, len(binary), 8):
        byte = binary[i : i + 8]
        text += chr(int(byte, 2))
    return text

# rsa修改
def binary_to_bytes(binary_list):
    # 将每个二进制字符串转换为整数
    byte_list = [int(binary_str, 2) for binary_str in binary_list]
    # 将整数列表转换为字节串
    rsa_bytes = bytes(byte_list)
    # 加载私钥  
    return rsa_bytes

def actor_dct(image_path, data_r):
    alpha = 30
    blocksize = 8
    watermark_text = data_r
    # rsa修改
    watermark_binary = bytes_to_binary(watermark_text)
    watermark = np.array([int(bit) for bit in watermark_binary], dtype=np.uint8).reshape((-1, len(watermark_binary) // 8))

    background = cv2.imread(image_path)
    yuv_background = cv2.cvtColor(background, cv2.COLOR_BGR2YUV)
    U = yuv_background[..., 1]

    dct_emb = DCT_Embed(background=U, watermark=watermark, block_size=blocksize, alpha=alpha)

    background_dct_blocks = dct_emb.dct_blkproc(background=U)
    embed_watermark_blocks = dct_emb.dct_embed(dct_data=background_dct_blocks, watermark=watermark)
    synthesis = dct_emb.idct_embed(dct_data=embed_watermark_blocks)
    print(synthesis)
    yuv_background[..., 1] = synthesis
    rgb_synthesis = cv2.cvtColor(yuv_background, cv2.COLOR_YUV2RGB)
    cv2.imwrite("ans_dct.png", cv2.cvtColor(rgb_synthesis, cv2.COLOR_RGB2BGR))
    # np.save("synthesis.npy", synthesis)  # 保存合成后的数据
    return watermark.shape


def reader_dct(image_path, len_data):
    block_size = 8
    alpha = 30
    watermark_size = (8, len_data)

    # 从 image_path 加载合成图像
    image = cv2.imread(image_path)
    yuv_image = cv2.cvtColor(image, cv2.COLOR_BGR2YUV)

    # 提取YUV图像的U通道
    synthesis = yuv_image[..., 1]
    print(synthesis)

    dct_emb = DCT_Embed(background=synthesis, watermark=np.zeros(watermark_size), block_size=block_size, alpha=alpha)
    np.random.seed(2024)
    dct_emb.k1 = np.random.randn(block_size)
    dct_emb.k2 = np.random.randn(block_size)

    extracted_watermark = dct_emb.dct_extract(synthesis=synthesis, watermark_size=watermark_size)
    extracted_watermark_binary = "".join(str(int(bit)) for bit in extracted_watermark.flatten())
    #rsa修改
    extracted_text = binary_to_bytes(extracted_watermark_binary)

    print(f"Extracted watermark (binary): {extracted_watermark_binary}")
    print(f"Extracted text: {extracted_text}")
    return extracted_text
